#!/usr/bin/env python3
"""Generate a specialized fullPipeline shell script based on the master
`fullPipeline.sh` template.

Usage:
    python generate_full_pipeline.py <gpu_id> <collect_num> <custom_name> <embodiment> <action_type> <trim_arm> <camera_use_full> <dp3_config_name>

The generated script will be named::

    fullPipeline-g<gpu_id>-num<collect_num>-<custom_name>-<embodiment>-<action_type>-ta<trim_arm>-cam<camera_use_full>.sh

and placed in the same directory as this Python script.
"""
import argparse
import os
import stat
import sys
import yaml
import shutil
from pathlib import Path

# Name of the master template relative to this script
MASTER_TEMPLATE = "fullPipeline.sh"


def read_master_template() -> list[str]:
    """Return the lines from the master template script."""
    template_path = Path(__file__).with_name(MASTER_TEMPLATE)
    if not template_path.exists():
        sys.exit(f"Error: master template '{MASTER_TEMPLATE}' not found next to this script.")
    return template_path.read_text(encoding="utf-8").splitlines()


def create_config_file(name: str, collect_num: int, embodiment: str, action_type: str, trim_arm: bool, camera_use_full: bool):
    """Create the specialized config file by copying and modifying the template."""
    # Source and destination paths
    source_path = Path(__file__).parent.parent/"RoboTwin/task_config/demo_clean.yml"
    dest_path = Path(__file__).parent.parent/f"RoboTwin/task_config/{name}.yml"
    
    if not source_path.exists():
        sys.exit(f"Error: source config file '{source_path}' not found.")
    
    # Copy the template file
    shutil.copy2(source_path, dest_path)
    
    # Read the YAML content
    with open(dest_path, 'r', encoding='utf-8') as f:
        content = yaml.safe_load(f)
    
    # Modify the content
    content['episode_num'] = collect_num
    content['action_type'] = action_type
    content['trim_arm'] = trim_arm
    
    # Update embodiment (replace all instances of franka-panda)
    if 'embodiment' in content and isinstance(content['embodiment'], list):
        content['embodiment'] = [embodiment, embodiment, 0.8]
    
    # Update pcd_camera_names based on camera_use_full
    if 'camera' in content and 'pcd_camera_names' in content['camera']:
        if camera_use_full:
            # Use full camera list
            content['camera']['pcd_camera_names'] = ['midBack_camera','midFront_camera','midLeft_camera','midRight_camera','midHead_camera']
        else:
            # Use only midBack_camera
            content['camera']['pcd_camera_names'] = ['midBack_camera']
    
    # Add zarr_path field - construct the path based on the naming convention
    # Pattern: place_shoe-{name}-{collect_num}-{action_type}.zarr
    zarr_filename = f"place_shoe-{name}-{collect_num}-{action_type}.zarr"
    zarr_path = f"{Path(__file__).parent.parent}/RoboTwin/policy/DP3/data/{zarr_filename}"
    content['zarr_path'] = zarr_path
    
    # Write the modified content back
    with open(dest_path, 'w', encoding='utf-8') as f:
        yaml.dump(content, f, default_flow_style=False, sort_keys=False)
    
    print(f"Created config file: {dest_path}")
    print(f"Added zarr_path: {zarr_path}")


def substitute_params(lines: list[str], gpu_id: int, collect_num: int, custom_name: str, dp3_config_name: str, embodiment: str, action_type: str, trim_arm: bool, camera_use_full: bool) -> list[str]:
    """Return a list of lines with parameters substituted.

    The first eight lines of the template (parameter assignments) are removed.
    The `name=` assignment line is regenerated with actual values.
    `$gpu_id`/`$collect_num` and their brace variants are replaced with fixed values.
    `$name` is kept as a variable for readability.
    """
    # Skip the first eight lines that assign $1, $2, $3, $4, $5, $6, $7, $8
    body = lines[8:]

    # Prepare output lines
    output: list[str] = [
        "#!/usr/bin/env bash",
        f"# Auto-generated by generate_full_pipeline.py for gpu_id={gpu_id}, collect_num={collect_num}, custom_name={custom_name}, dp3_config_name={dp3_config_name}, embodiment={embodiment}, action_type={action_type}, trim_arm={trim_arm}, camera_use_full={camera_use_full}",
        "set -e",
        "",
    ]

    # Build the canonical name value once, including all parameters
    name_value = f"demo_clean_{custom_name}-{collect_num}-{embodiment}-{action_type}-ta{1 if trim_arm else 0}-cam{1 if camera_use_full else 0}"

    for line in body:
        stripped = line.strip()
        # Replace the `name=` line entirely
        if stripped.startswith("name="):
            output.append(f"name={name_value}")
            continue

        # Perform simple substitutions
        line = (
            line.replace("${gpu_id}", str(gpu_id))
                .replace("$gpu_id", str(gpu_id))
                .replace("${collect_num}", str(collect_num))
                .replace("$collect_num", str(collect_num))
                .replace("${custom_name}", custom_name)
                .replace("$custom_name", custom_name)
                .replace("${dp3_config_name}", dp3_config_name)
                .replace("$dp3_config_name", dp3_config_name)
                .replace("${embodiment}", embodiment)
                .replace("$embodiment", embodiment)
                .replace("${action_type}", action_type)
                .replace("$action_type", action_type)
                .replace("${name}", name_value)
                .replace("$name", name_value)
                .replace("${root_path}", os.path.join(str(Path(__file__).parent.parent), "RoboTwin"))
        )
        output.append(line)

    output.append("")  # Ensure final newline
    return output


def build_script(gpu_id: int, collect_num: int, custom_name: str, embodiment: str, action_type: str, trim_arm: bool, camera_use_full: bool, dp3_config_name: str) -> str:
    """Create the full text of the specialized script."""
    master_lines = read_master_template()
    specialized_lines = substitute_params(master_lines, gpu_id, collect_num, custom_name, dp3_config_name, embodiment, action_type, trim_arm, camera_use_full)
    return "\n".join(specialized_lines)


def main():
    parser = argparse.ArgumentParser(description="Generate specialized fullPipeline shell script.")
    parser.add_argument("gpu_id", type=int, help="GPU ID to use")
    parser.add_argument("collect_num", type=int, help="Number of episodes to collect")
    parser.add_argument("custom_name", help="Custom name tag")
    parser.add_argument("embodiment", help="Embodiment name (e.g., franka-panda)")
    parser.add_argument("action_type", help="Action type")
    parser.add_argument("trim_arm", type=int, choices=[0, 1], help="Trim arm setting (0 or 1)")
    parser.add_argument("camera_use_full", type=int, choices=[0, 1], help="Use full camera list (0 or 1)")
    parser.add_argument("--dp3_config_name", help="DP3 config name", default="robot_dp3-7212-mlp2512-fff-s")
    args = parser.parse_args()

    # Assign variables in new order
    gpu_id = args.gpu_id
    collect_num = args.collect_num
    custom_name = args.custom_name
    embodiment = args.embodiment
    action_type = args.action_type
    trim_arm = bool(args.trim_arm)
    camera_use_full = bool(args.camera_use_full)
    dp3_config_name = args.dp3_config_name

    script_name = f"fullPipeline-g{gpu_id}-num{collect_num}-{custom_name}-{embodiment}-{action_type}-ta{1 if trim_arm else 0}-cam{1 if camera_use_full else 0}.sh"
    script_path = Path(__file__).with_name(script_name)

    # Create the config file first
    name_value = f"demo_clean_{custom_name}-{collect_num}-{embodiment}-{action_type}-ta{1 if trim_arm else 0}-cam{1 if camera_use_full else 0}"
    create_config_file(name_value, collect_num, embodiment, action_type, trim_arm, camera_use_full)

    content = build_script(gpu_id, collect_num, custom_name, embodiment, action_type, trim_arm, camera_use_full, dp3_config_name)

    script_path.write_text(content, encoding="utf-8")

    # Make the script executable
    mode = script_path.stat().st_mode
    script_path.chmod(mode | stat.S_IEXEC)

    print(f"Generated {script_path}")


if __name__ == "__main__":
    main() 